{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "exempt-legislation",
   "metadata": {
    "id": "exempt-legislation"
   },
   "source": [
    "<center>\n",
    "<h4>CDS 110, Lecture 6b</h4>\n",
    "<font color=blue><h1>Trajectory Tracking for a Kinematic Car</h1></font>\n",
    "<h3>Richard M. Murray, Winter 2024</h3>\n",
    "</center>\n",
    "\n",
    "[Open in Google Colab](https://colab.research.google.com/drive/12VSFMqM6HVyj8TY_3zb0AnsJrG6UeLKF)\n",
    "\n",
    "This notebook contains an example of using trajectory tracking for a (nonlinear) state space system.  The controller is of the form\n",
    "\n",
    "$$\n",
    "  u = u_\\text{d} − K (x − x_\\text{d}),\n",
    "$$\n",
    "\n",
    "where $x_\\text{d}, u_\\text{d}$ is a feasible trajectory, and $K$ is a feedback gain first computed around a nominal condition and then computed using gain scheduling."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "corresponding-convenience",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import the packages needed for the examples included in this notebook\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import itertools\n",
    "from cmath import sqrt\n",
    "from math import pi\n",
    "try:\n",
    "  import control as ct\n",
    "  print(\"python-control\", ct.__version__)\n",
    "except ImportError:\n",
    "  !pip install control\n",
    "  import control as ct\n",
    "import control.optimal as opt\n",
    "import control.flatsys as fs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "corporate-sense",
   "metadata": {
    "id": "corporate-sense"
   },
   "source": [
    "## Vehicle Steering Dynamics\n",
    "\n",
    "The vehicle dynamics are given by a simple bicycle model:\n",
    "\n",
    "<table>\n",
    "<tr>\n",
    "    <td width=\"50%\"><img src=\"https://fbswiki.org/wiki/images/5/52/Kincar.png\" width=480></td>\n",
    "    <td width=\"50%\">\n",
    "$$\\large\n",
    "\\begin{aligned}\n",
    "  \\dot x &= \\cos\\theta\\, v \\\\\n",
    "  \\dot y &= \\sin\\theta\\, v \\\\\n",
    "  \\dot\\theta &= \\frac{v}{l} \\tan \\delta\n",
    "\\end{aligned}\n",
    "$$\n",
    "    </td>\n",
    "</tr>\n",
    "</table>\n",
    "\n",
    "We take the state of the system as $(x, y, \\theta)$ where $(x, y)$ is the position of the vehicle in the plane and $\\theta$ is the angle of the vehicle with respect to horizontal.  The vehicle input is given by $(v, \\delta)$ where $v$ is the forward velocity of the vehicle and $\\delta$ is the angle of the steering wheel.  The model includes saturation of the vehicle steering angle."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "naval-pizza",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Code to model vehicle steering dynamics\n",
    "\n",
    "# Function to compute the RHS of the system dynamics\n",
    "def kincar_update(t, x, u, params):\n",
    "    # Get the parameters for the model\n",
    "    l = params['wheelbase']             # vehicle wheelbase\n",
    "    deltamax = params['maxsteer']         # max steering angle (rad)\n",
    "\n",
    "    # Saturate the steering input\n",
    "    delta = np.clip(u[1], -deltamax, deltamax)\n",
    "\n",
    "    # Return the derivative of the state\n",
    "    return np.array([\n",
    "        np.cos(x[2]) * u[0],            # xdot = cos(theta) v\n",
    "        np.sin(x[2]) * u[0],            # ydot = sin(theta) v\n",
    "        (u[0] / l) * np.tan(delta)      # thdot = v/l tan(delta)\n",
    "    ])\n",
    "\n",
    "kincar_params={'wheelbase': 3, 'maxsteer': 0.5}\n",
    "\n",
    "# Create nonlinear input/output system\n",
    "kincar = ct.nlsys(\n",
    "    kincar_update, None, name=\"kincar\", params=kincar_params,\n",
    "    inputs=('v', 'delta'), outputs=('x', 'y', 'theta'),\n",
    "    states=('x', 'y', 'theta'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6340dbd4-7867-47ad-aefb-1bea7f6ad566",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Utility function to plot lane change manuever\n",
    "def plot_lanechange(t, y, u, figure=None, yf=None, label=None):\n",
    "    # Plot the xy trajectory\n",
    "    plt.subplot(3, 1, 1, label='xy')\n",
    "    plt.plot(y[0], y[1], label=label)\n",
    "    plt.xlabel(\"x [m]\")\n",
    "    plt.ylabel(\"y [m]\")\n",
    "    if yf is not None:\n",
    "        plt.plot(yf[0], yf[1], 'ro')\n",
    "\n",
    "    # Plot x and y as functions of time\n",
    "    plt.subplot(3, 2, 3, label='x')\n",
    "    plt.plot(t, y[0])\n",
    "    plt.ylabel(\"$x$ [m]\")\n",
    "\n",
    "    plt.subplot(3, 2, 4, label='y')\n",
    "    plt.plot(t, y[1])\n",
    "    plt.ylabel(\"$y$ [m]\")\n",
    "\n",
    "    # Plot the inputs as a function of time\n",
    "    plt.subplot(3, 2, 5, label='v')\n",
    "    plt.plot(t, u[0])\n",
    "    plt.xlabel(\"Time $t$ [sec]\")\n",
    "    plt.ylabel(\"$v$ [m/s]\")\n",
    "\n",
    "    plt.subplot(3, 2, 6, label='delta')\n",
    "    plt.plot(t, u[1])\n",
    "    plt.xlabel(\"Time $t$ [sec]\")\n",
    "    plt.ylabel(\"$\\\\delta$ [rad]\")\n",
    "\n",
    "    plt.subplot(3, 1, 1)\n",
    "    plt.title(\"Lane change manuever\")\n",
    "    if label:\n",
    "        plt.legend()\n",
    "    plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "BAsKLMWWK3W2",
   "metadata": {
    "id": "BAsKLMWWK3W2"
   },
   "source": [
    "## State feedback controller\n",
    "\n",
    "We start by designing a state feedback controller that can be used to stabilize the system.  We design the controller around a nominal forward speed of 10 m/s and then apply this to the vehicle at different speeds."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "g7DztIjmK2K_",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute the linearization of the dynamics at a nominal point\n",
    "x_nom = np.array([0, 0, 0])\n",
    "u_nom = np.array([5, 0])\n",
    "P = ct.linearize(kincar, x_nom, u_nom)    # Linearized systems\n",
    "print(P)\n",
    "\n",
    "Qx = np.diag([1, 10, 0.1])\n",
    "Qu = np.diag([1, 1])\n",
    "K, _, _ = ct.lqr(P.A, P.B, Qx, Qu)\n",
    "print(K)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "szvKKh6rLgkt",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the closed loop system using create_statefbk_iosystem\n",
    "?ct.create_statefbk_iosystem\n",
    "ctrl, clsys = ct.create_statefbk_iosystem(\n",
    "    kincar, K, xd_labels=['xd', 'yd', 'thetad'], ud_labels=['vd', 'deltad'])\n",
    "print(clsys)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "gow-ZEerMCw7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a trajectory corresponding to a slow lane change\n",
    "x0 = np.array([0, -2, 0]); u0 = [10, 0]\n",
    "xf = np.array([100, 2, 0])\n",
    "Tf = 10\n",
    "timepts = np.linspace(0, Tf, 20)\n",
    "\n",
    "straight_line = (               # straight line from start to end with nominal input\n",
    "    np.array([x0 + (xf - x0) * t/Tf for t in timepts]).transpose(),\n",
    "    u0\n",
    ")\n",
    "\n",
    "desired = opt.solve_ocp(\n",
    "    kincar, timepts, x0,\n",
    "    cost=opt.quadratic_cost(kincar, None, Qu, u0=u0),\n",
    "    terminal_constraints=opt.state_range_constraint(kincar, xf, xf),\n",
    "    initial_guess=straight_line)\n",
    "\n",
    "plot_lanechange(desired.time, desired.states, desired.inputs, yf=xf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "NLa4dbI8PWhY",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simulate the system with an initial condition error\n",
    "# Use t_eval to evaluate at points between inputs\n",
    "actual = ct.input_output_response(\n",
    "    clsys, timepts, [desired.states, desired.inputs],\n",
    "    X0=[-3, -5, 0], t_eval=np.linspace(0, Tf, 500))\n",
    "\n",
    "plot_lanechange(actual.time, actual.states, actual.outputs[3:])\n",
    "plot_lanechange(desired.time, desired.states, desired.inputs, yf=xf)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "TKyc2jOiWJBe",
   "metadata": {
    "id": "TKyc2jOiWJBe"
   },
   "source": [
    "Note that the value of $\\delta$ is very large at the start.  This is truncated in the model so that it does not exceed $\\pm 0.5$ rad."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c6c4b9b",
   "metadata": {
    "id": "6c6c4b9b"
   },
   "source": [
    "## Reference trajectory subsystem\n",
    "\n",
    "In addition to generating a trajectory for the system, we can also create $x_\\text{d}$ and $u_\\text{d}$ corresponding to reference inputs $r_y$ and $r_v$.\n",
    "\n",
    "The reference trajectory block below generates a simple trajectory for the system given the desired speed (vref) and lateral position (yref).  The trajectory consists of a straight line of the form (vref * t, yref, 0) with nominal\n",
    "input (vref, 0)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "significant-november",
   "metadata": {},
   "outputs": [],
   "source": [
    "# System state: none\n",
    "# System input: vref, yref\n",
    "# System output: xd, yd, thetad, vd, deltad\n",
    "# System parameters: none\n",
    "#\n",
    "def trajgen_output(t, x, u, params):\n",
    "    vref, yref = u\n",
    "    return np.array([vref * t, yref, 0, vref, 0])\n",
    "\n",
    "# Define the trajectory generator as an input/output system\n",
    "trajgen = ct.nlsys(\n",
    "    None, trajgen_output, name='trajgen',\n",
    "    inputs=('vref', 'yref'),\n",
    "    outputs=('xd', 'yd', 'thetad', 'vd', 'deltad'))\n",
    "\n",
    "print(trajgen)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0w5s56uUWw-v",
   "metadata": {
    "id": "0w5s56uUWw-v"
   },
   "source": [
    "## Step responses\n",
    "\n",
    "To explore the dynamics of the system, we create a set of lane changes at different forward speeds.  Since the linearization depends on the speed, this means that the closed loop performance of the system will vary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "mtGLwMQkXEzw",
   "metadata": {},
   "outputs": [],
   "source": [
    "steering_fixed = ct.interconnect(\n",
    "    [kincar, ctrl, trajgen],\n",
    "    inputs=['vref', 'yref'],\n",
    "    outputs=kincar.output_labels + kincar.input_labels\n",
    ")\n",
    "print(steering_fixed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "sz7NaJTGXua1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up the simulation conditions\n",
    "yref = 1\n",
    "T = np.linspace(0, 5, 100)\n",
    "\n",
    "# Do an iteration through different speeds\n",
    "for vref in [2, 5, 20]:\n",
    "    # Simulate the closed loop controller response\n",
    "    tout, yout = ct.input_output_response(\n",
    "        steering_fixed, T, [vref * np.ones(len(T)), yref * np.ones(len(T))],\n",
    "        params={'maxsteer': 1})\n",
    "\n",
    "    # Plot the results\n",
    "    plot_lanechange(tout, yout, yout[3:])\n",
    "\n",
    "# Label the different curves\n",
    "plt.subplot(3, 1, 1)\n",
    "plt.legend([\"$v_d$ = \" + f\"{vref}\" for vref in [2, 10, 20]])\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3cc26675",
   "metadata": {
    "id": "3cc26675"
   },
   "source": [
    "## Gain scheduled controller\n",
    "\n",
    "For this system we use a simple schedule on the forward vehicle velocity and\n",
    "place the poles of the system at fixed values.  The controller takes the\n",
    "current and desired vehicle position and orientation plus the velocity\n",
    "velocity as inputs, and returns the velocity and steering commands.\n",
    "\n",
    "Linearizing the system about the desired trajectory, we obtain\n",
    "\n",
    "$$\n",
    "  \\begin{aligned}\n",
    "    A(x_\\text{d}) &= \\left. \\frac{\\partial f}{\\partial x} \\right|_{(x_\\text{d}, u_\\text{d})}\n",
    "      = \\left.\n",
    "        \\begin{bmatrix}\n",
    "          0 & 0 & -\\sin\\theta_\\text{d}\\, v_\\text{d} \\\\ 0 & 0 & \\cos\\theta_\\text{d}\\, v_\\text{d} \\\\ 0 & 0 & 0\n",
    "        \\end{bmatrix}\n",
    "        \\right|_{(x_\\text{d}, u_\\text{d})}\n",
    "      = \\begin{bmatrix}\n",
    "          0 & 0 & 0 \\\\ 0 & 0 & v_\\text{d} \\\\ 0 & 0 & 0\n",
    "         \\end{bmatrix}, \\\\\n",
    "    B(x_\\text{d}) &= \\left. \\frac{\\partial f}{\\partial u} \\right|_{(x_\\text{d}, u_\\text{d})}\n",
    "     = \\begin{bmatrix}\n",
    "       1 & 0 \\\\ 0 & 0 \\\\ 0 & v_\\text{d}/l\n",
    "       \\end{bmatrix}.\n",
    "  \\end{aligned}\n",
    "$$\n",
    "\n",
    "We see that these matrices depend only on $\\theta_\\text{d}$ and $v_\\text{d}$, so we choose these as the scheduling variables and design a controller of the form\n",
    "\n",
    "$$\n",
    "u = u_\\text{d} - K(\\mu) (x - x_\\text{d})\n",
    "$$\n",
    "\n",
    "where $\\mu = (\\theta_\\text{d}, v_\\text{d})$ and we interpolate the gains based on LQR controllers computed at a fixed set of points $\\mu_i$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "another-milwaukee",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the points for the scheduling variables\n",
    "gs_speeds = [2, 10, 20]\n",
    "gs_angles = np.linspace(-pi, pi, 4)\n",
    "\n",
    "# Create controllers at each scheduling point (\n",
    "points = [np.array([speed, angle])\n",
    "          for speed in gs_speeds for angle in gs_angles]\n",
    "gains = [np.array(ct.lqr(kincar.linearize(\n",
    "    [0, 0, angle], [speed, 0]), Qx, Qu)[0])\n",
    "    for speed in gs_speeds for angle in gs_angles]\n",
    "print(f\"{points=}\")\n",
    "print(f\"{gains=}\")\n",
    "\n",
    "# Create the gain scheduled system\n",
    "ctrl_gs, _ = ct.create_statefbk_iosystem(\n",
    "    kincar, (gains, points), name='controller',\n",
    "    xd_labels=['xd', 'yd', 'thetad'], ud_labels=['vd', 'deltad'],\n",
    "    gainsched_indices=['vd', 'theta'], gainsched_method='linear')\n",
    "print(ctrl_gs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ca5ab53",
   "metadata": {
    "id": "4ca5ab53"
   },
   "source": [
    "## System construction\n",
    "\n",
    "The input to the full closed loop system is the desired lateral position and the desired forward velocity.  The output for the system is taken as the full vehicle state plus the velocity of the vehicle.\n",
    "\n",
    "We construct the system using the `ct.interconnect` function and use signal labels to keep track of everything.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "editorial-satisfaction",
   "metadata": {},
   "outputs": [],
   "source": [
    "steering_gainsched = ct.interconnect(\n",
    "    [trajgen, ctrl_gs, kincar], name='steering',\n",
    "    inputs=['vref', 'yref'],\n",
    "    outputs=kincar.output_labels + kincar.input_labels\n",
    ")\n",
    "print(steering_gainsched)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47f5d528",
   "metadata": {
    "id": "47f5d528"
   },
   "source": [
    "## System simulation\n",
    "\n",
    "We now simulate the gain scheduled controller for a step input in the $y$ position, using a range of vehicle speeds $v_\\text{d}$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "smoking-trail",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the reference trajectory for the y position\n",
    "# plt.plot([0, 5], [yref, yref], 'k-', linewidth=0.6)\n",
    "\n",
    "# Find the signals we want to plot\n",
    "y_index = steering_gainsched.find_output('y')\n",
    "v_index = steering_gainsched.find_output('v')\n",
    "\n",
    "# Do an iteration through different speeds\n",
    "for vref in [2, 5, 20]:\n",
    "    # Simulate the closed loop controller response\n",
    "    tout, yout = ct.input_output_response(\n",
    "        steering_gainsched, T, [vref * np.ones(len(T)), yref * np.ones(len(T))],\n",
    "        X0=[0, 0, 0], params={'maxsteer': 0.5}\n",
    "    )\n",
    "\n",
    "    # Plot the results\n",
    "    plot_lanechange(tout, yout, yout[3:])\n",
    "\n",
    "# Label the different curves\n",
    "plt.subplot(3, 1, 1)\n",
    "plt.legend([\"$v_d$ = \" + f\"{vref}\" for vref in [2, 10, 20]])\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f571b2b",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
